"""Local worker implementation for executing tasks on the local machine.

This module contains a worker implementation that runs the worker.py script
as a local subprocess.
"""

import asyncio
import os
import sys
from typing import Optional

from loguru import logger

from tasks.model import TaskStatus
from tasks.worker.constants import EXIT_SUCCESS, EXIT_FAILURE
from tasks.worker.interface import WorkerInterface, WorkspaceSpec


class LocalWorker(WorkerInterface):
    """Worker implementation that runs tasks as local subprocesses."""
    
    def __init__(self, workspace: WorkspaceSpec):
        """Initialize a local worker.
        
        Args:
            workspace: The workspace specification for the task
        """
        self.workspace = workspace
        self.process: Optional[asyncio.subprocess.Process] = None
        self._status = TaskStatus.PENDING
    
    async def start(self) -> None:
        """Start the task execution as a local subprocess.
        
        Returns:
            None
            
        Raises:
            RuntimeError: If the task is already running
        """
        if self.process is not None:
            raise RuntimeError("Task is already running")
        
        self._status = TaskStatus.RUNNING
        
        cwd = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        env = os.environ.copy()
        env["PYTHONPATH"] = os.path.dirname(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        )
        
        self.process = await asyncio.create_subprocess_exec(
            "python",
            "-m",
            "tasks.worker.worker",
            "run",
            self.workspace.task_dir,
            cwd=cwd,
            env=env,
        )
        
        logger.info(f"Started worker process for task {self.workspace.task_hash}")
    
    async def get_status(self) -> TaskStatus:
        """Get the current status of the task.
        
        Returns:
            TaskStatus: The current status of the task
        """
        if self.process is None:
            return TaskStatus.PENDING
        
        if self.process.returncode is None:
            return TaskStatus.RUNNING
        
        if self.process.returncode == EXIT_SUCCESS:
            return TaskStatus.COMPLETED
        else:
            return TaskStatus.FAILED
    
    async def wait(self) -> int:
        """Wait for the task to complete and return the exit code.
        
        Returns:
            int: Exit code (0 for success, non-zero for failure)
            
        Raises:
            RuntimeError: If the task has not been started
        """
        if self.process is None:
            raise RuntimeError("Task has not been started")
        
        return await self.process.wait()
    
    async def cancel(self) -> None:
        """Cancel the running task.
        
        Returns:
            None
        """
        if self.process is not None and self.process.returncode is None:
            self.process.kill()
            self._status = TaskStatus.CANCELLED
            logger.info(f"Cancelled worker process for task {self.workspace.task_hash}")
    
    @property
    def returncode(self) -> Optional[int]:
        """Get the return code of the task if it has completed.
        
        Returns:
            Optional[int]: Return code if task has completed, None otherwise
        """
        if self.process is None:
            return None
        
        return self.process.returncode 